import cv2
import time
import os
import json
from LipFeaturaExtract.landmarks import FaceLandmarks,normalization
import numpy as np

# 数据集路径
lipdata_path = "/home/lisen/tool/dataset/唇语数据集/lipread_mp4"
save_dir = "/home/lisen/tool/PyProjects/唇语识别/dataset/dataset_word/dataset_v2"
pattern="train"

if not os.path.exists(save_dir):
    os.mkdir(save_dir)
extract_words=[]
if not os.path.exists(os.path.join(save_dir,pattern)):
    os.mkdir(os.path.join(save_dir,pattern))
else:
    word_files=os.listdir(os.path.join(save_dir,pattern))
    for word_name in word_files:
        extract_words.append(word_name.rstrip(".json"))

# 生成脸部特征提取实例
faceLandmarks=FaceLandmarks(scaling_factor=1,extract_face=False)

# 提取文件下所有唇部特征并保存
words_list = os.listdir(lipdata_path)
word_count = 0  # 单词计数
if pattern=="train":
    # words_list=[word for word in words_list if word not in ["RIGHT","BEFORE","RIGHTS","AFTER"]]
    # words_list=words_list[:46]+["RIGHT","BEFORE","RIGHTS","AFTER"]
    words_list=["RIGHT","BEFORE","RIGHTS","AFTER"]
else:
    words_list=os.listdir(os.path.join(save_dir,"train"))
    words_list=[word.strip(".json") for word in words_list]
    # words_list = ["RIGHT", "BEFORE", "RIGHTS", "AFTER"]

if not os.path.exists(os.path.join(save_dir, "ids.json")):
    label2id={word:index for index,word in enumerate(words_list)}
    id2label=dict(zip(label2id.values(), label2id.keys()))
    print(id2label)
    exit()
    with open(os.path.join(save_dir, "ids.json"), "w", encoding="utf-8") as f:
        json.dump({'label2id': label2id,
                   'id2label':id2label},
                  fp=f, ensure_ascii=False)

words_list_len = len(words_list)
for word in words_list:
    word_count += 1
    if word not in extract_words:
        word_lip_data=[]
        print("========Extract word:{} ({}/{})==========".format(
            word, word_count,words_list_len))

        lip_word_pattern = os.path.join(word, pattern)
        files = os.listdir(os.path.join(lipdata_path, lip_word_pattern))
        video_files = [file for file in files if ".mp4" in file]  # 提取视频流文件路径

        # 提取特征
        video_files_len = len(video_files)  # 每个词的个数
        video_conut=0   # 视频计数
        for video in video_files:
            count = 0
            cap = cv2.VideoCapture(os.path.join(lipdata_path, lip_word_pattern, video))  # 打开摄像头
            ret = True
            video_lip_data = []
            start_time = time.clock()  # 开始时间
            while ret:
                ret, frame = cap.read()  # 读取一帧图片
                if ret:
                    # start = time.clock()  # 开始时间
                    # face_locations,lip_features=faceLandmarks.extractLandmarks(frame)
                    lip_features = faceLandmarks.extractLandmarks(frame)

                    # # 绘制特征区域
                    # for index,feature in enumerate(lip_features):
                    #     # print(index+face_locations[index][1])
                    #     for point in feature:
                    #         cv2.circle(frame, (point[0]+face_locations[index][3],
                    #                            point[1]+face_locations[index][0]),
                    #                    1, (0, 0, 255), -1)
                    # for face_location in face_locations:
                    #     cv2.rectangle(frame,(face_location[3],face_location[0]),
                    #                   (face_location[1],face_location[2]),
                    #                   (0,0,255),1)
                    # 对数据进行归一化
                    if len(lip_features)==1:
                        w=frame.shape[1]
                        h=frame.shape[0]
                        # h=face_locations[0][2]-face_locations[0][0]
                        # w=face_locations[0][1]-face_locations[0][3]
                        lip_features = normalization(lip_features,w,h)[0]
                        video_lip_data.append(lip_features)
                        count+=1

                    # end = time.clock()
                    # 显示图像
                    # cv2.imshow("src", frame)
                    # print("帧数：{:.3f}fps".format(1 / (end - start)))
                    # key = cv2.waitKey()
                    # if key == 27:
                    #     ret = False
            word_lip_data.append(video_lip_data)
            video_conut+=1
            end_time = time.clock()
            print("schedule:{}/{}  time:{:.6f}".format(video_conut, video_files_len,end_time-start_time))
            cap.release()

        # 使用json保存数据
        with open(os.path.join(save_dir, pattern, word + ".json"), "w", encoding="utf-8") as f:
            json.dump({'lip_features': word_lip_data},
                      fp=f, ensure_ascii=False)
        print("============Extract one word!=================")
print("提取完成")


